#include "LSR.h"

/**
*  Basic contructor. Called once most of the information has already been discovered
*  and collects it into an LSR object. Calls findLine to get line points. 'keeper' property
*  flags whether the LSR is worth keeping or not.
*/
LSR::LSR(int area, int *cells, double orientation, Mat img){
    if(area < MIN_CELLS) {
        keeper = 0;
        orientation = -5;
        length = -1;
    }
    else{
        int totLength = XDIM*YDIM;
        this->area = area;
        this->cells = cells;
        this->orientation = orientation;

        int x_sum = 0;
        int y_sum = 0;
        int cellAddr;
        int *cellCoords;
        int x_min = XDIM;    
        int y_min = YDIM;
        int x_max = -1;
        int y_max = -1;
        int *boundingBox = new int[4];
        for(int i=0; i<area; i++){
            cellAddr = cells[i];
            cellCoords = addrToCoords(cellAddr);
            if(cellCoords[0]<x_min){
                boundingBox[0]=cellCoords[0];
                x_min = cellCoords[0];
            }
            if(cellCoords[1]<y_min){
                boundingBox[1]=cellCoords[1];
                y_min = cellCoords[1];
            }
            if(cellCoords[0]>x_max){
                boundingBox[2]=cellCoords[0];
                x_max = cellCoords[0];
            }
            if(cellCoords[1]>y_max){
                boundingBox[3]=cellCoords[1];
                y_max = cellCoords[1];
            }
            x_sum += cellCoords[0];
            y_sum += cellCoords[1];
        }
        int x_avg = x_sum/area;
        int y_avg = y_sum/area;
        centroid = coordsToAddr(x_avg, y_avg);
        linePoints = findLine(boundingBox, centroid, orientation);
        depths = new double[3];
        depths[0] = depthAtPoint(linePoints[0], img);
        depths[1] = depthAtPoint(centroid, img);
        depths[2] = depthAtPoint(linePoints[1], img);
        length = pointDist(linePoints[0], linePoints[1]);
        if(length < MIN_LENGTH || length > MAX_LENGTH)
            keeper = 0; //NOTE: Angle bias could be introduced here.
        else keeper = 1;
    }
}

int *LSR::getCells(){
    return cells;
}

int LSR::getArea(){
    return area;
}

int LSR::getCentroid(){
    return centroid;
}

double LSR::getOrientation(){
    return orientation;
}

double LSR::getLength(){
    return length;
}

bool LSR::isKeeper(){
    return keeper;
}

int *LSR::getLinePoints(){
    return linePoints;
}

double *LSR::getDepths(){
    return depths;
}

/**
* Finds the depth at a given point by taking the darkest (= closest) pixel 
* in a 2-pixel radius. TODO: Make sure this holds for the optical-flow depth images.
*/
double depthAtPoint(int addr, Mat img){
    double darkest = 256.0;
    int *coords = addrToCoords(addr);
    int xCoord = coords[0];
    int yCoord = coords[1];
    for(int j = yCoord-2; j <= yCoord+2; j++){
        if(j<0 || j>YDIM)
            continue;
        const double* row_j = img.ptr<double>(j);
        for(int i = xCoord-2; i <= xCoord+2; i++){
            if(i<0 || i>XDIM)
                continue;
            double value = row_j[i];
            //cout << "value at " << i<< ", " << j << " is: " << value << endl;
            if(value < darkest)
                darkest = value;
        }
    }
    //cout << "At addr " << addr << ", darkest is: " << darkest << endl;
    return darkest;
}

/**
*  This function takes the magnitudes and list of visited cells and returns the biggest 
*  cell addr that hasn't been visited. Returns -1 if we're all out. (Not too efficient)
*/
int getNextPix(int* mags, bool* used, int length){
    int max = -1;
    int maxInd = -1;
    for(int i=0; i < length; i++){
        if(mags[i]>max && used[i] == 0){  //biggest seen and not yet taken
            max = mags[i];
            maxInd = i;
        }
    }
    return maxInd;
}

/** 
*  This is based on the pseudocode in the von Gioi article. Takes a seed cell and builds an LSR.
*  Returns number of cells.
*/
double* regionGrow(int *region, double *angles, int *mags, bool *used, Mat drawImg){
    int seed = region[0];
    used[seed] = 1;
    int cells = 1;
    double regionTheta = angles[seed];
    double Sx = cos(regionTheta)*mags[seed];
    double Sy = sin(regionTheta)*mags[seed];
    int nChecked = 0; //number of cells in region whose neighbors have been checked
    while(nChecked < cells){			//Check this part
        int checkCell = region[nChecked];
        //int *neighbors = findAllNeighbors(checkCell);
        int neighborhood_area = 4*NEIGHBORHOOD_RADIUS*(NEIGHBORHOOD_RADIUS + 1);
        int *neighbors = findNeighborhood(checkCell, NEIGHBORHOOD_RADIUS);
        for(int i=0; i<neighborhood_area; i++){
            int thisNeighbor = neighbors[i];
            if(thisNeighbor != -1){
		double thisAngle = angles[thisNeighbor];
                if(thisAngle != -5 && used[thisNeighbor] == 0){
                double thetaMax = regionTheta + ANGLE_TOL;
                double thetaMin = regionTheta - ANGLE_TOL;
                //I need to separate out the cases where the thetaMax/Min are outside our bounds
		bool inRange = 0;
		bool edgeCase = 0;
		if(thetaMax > PI) {
		    thetaMax -= PI;             
		    edgeCase = 1;
		}
		if(thetaMin <= 0.001) {
		    thetaMin += PI;	     
		    edgeCase = 1;
		}
                // In edge cases, the angle in question needs to only be < Max OR > Min
                if(edgeCase){
                    if(thisAngle <= thetaMax || thisAngle >= thetaMin) inRange = 1;
		} else { // Non-edge cases must be both.
		    if(thisAngle <= thetaMax && thisAngle >= thetaMin) inRange = 1;
		}
                if(inRange){
                    region[cells] = thisNeighbor;
                    used[thisNeighbor] = 1;
		    if(SHOULD_DRAW){
			int *pt_coords = addrToCoords(thisNeighbor);
	    	        Point pt1 = Point(pt_coords[0],pt_coords[1]);
		        line(drawImg, pt1, pt1, Scalar(0,255,255), 1);
		    }
		    if(abs(regionTheta - thisAngle) > ANGLE_TOL){
			thisAngle += (regionTheta > thisAngle) ? PI : (-1*PI);
		    }
                    Sx += cos(thisAngle)*mags[thisNeighbor];
                    Sy += sin(thisAngle)*mags[thisNeighbor];
                    double newTheta = atan2(Sy, Sx);
		    if(newTheta > PI) regionTheta = newTheta - PI;
		    else if(newTheta <= 0) regionTheta = newTheta + PI;
		    else regionTheta = newTheta;
                    cells++;
                }
                }
            }
        }
        nChecked++;
    }
    double* stats = new double[2];
    stats[0] = cells;
    stats[1] = regionTheta;
    return stats;
}

/**
*  This takes a gradient angle and returns the level-line angle (perpendicular
*  to grad). It constrains answers to +, so that -3/4*pi \equiv 1/4*pi.
*    Rememer that the input and output are both in radians.
*/
double levelLine(double gradAngle){
    double retAng = gradAngle + (PI/2);
    if(retAng>PI) retAng -= PI;
    if(retAng <= 0.001) retAng += PI;
    return retAng;
}

/**
*  Populates an array of LSRs (Line Support Regions) given a starting image and blank
*  LSR array. Calls on regionGrow. 
*/
int findLSRs(LSR *allLSRs, Mat img){
    Mat img_gray, img_g, img_gx, img_gy, abs_img_gx, abs_img_gy;
    int totLength = XDIM*YDIM;
    Mat drawImg = img.clone();

    for(int i=0; i<SMOOTHING_ITERS; i++)
        GaussianBlur(img,img,Size(3,3),0,0,BORDER_DEFAULT);
    //imshow("Blurred", img);
    //waitKey(0);
    cvtColor(img,img_gray, CV_RGB2GRAY);
    Sobel(img_gray, img_gx, DDEPTH, 1, 0, 3, SCALE, DELTA, BORDER_DEFAULT); //gradient x
    Sobel(img_gray, img_gy, DDEPTH, 0, 1, 3, SCALE, DELTA, BORDER_DEFAULT); //gradient y
    convertScaleAbs(img_gx, abs_img_gx);
    convertScaleAbs(img_gy, abs_img_gy);
    addWeighted(abs_img_gx, 0.5, abs_img_gy,0.5,0,img_g);
    //imshow("Growing region...", img_g);
    //waitKey(0);

    // For processiong purposes:
    img_g.convertTo(img_g, CV_64F);
    img_gx.convertTo(img_gx, CV_64F);
    img_gy.convertTo(img_gy, CV_64F);
    img_gray.convertTo(img_gray, CV_64F);

    // Set up magnitudes, angles.
    int *mags = new int[totLength];
    double *angs = new double[totLength];
    for(int i = 0; i < YDIM; i++)
    {
        const double* Gi = img_g.ptr<double>(i);
        const double* Xi = img_gx.ptr<double>(i);
        const double* Yi = img_gy.ptr<double>(i);
        for(int j = 0; j < XDIM; j++){
	    int arrayAddr = coordsToAddr(j,i);
	    double val = Gi[j];
	    if(val>MAG_THRESH){
	        mags[arrayAddr] = (int)val;
		double gradAng = atan2(Yi[j],-1*Xi[j]);
	        angs[arrayAddr] = levelLine(gradAng);
	    } else {
	        mags[arrayAddr] = -1;
	        angs[arrayAddr] = -5; //Actual values are between 0 and PI
	    }
        }
    }
    int LSRcount = 0;
    bool *used = new bool[totLength];
    fill_n(used,totLength,0);
    int nextPix;
    for(int i=0; i<totLength; i++){ 
        nextPix = getNextPix(mags, used, totLength);
        if(nextPix == -1) {
            break;
        }
        int *region = new int[totLength/MIN_CELLS];
	region[0] = nextPix;
        double* regStats = regionGrow(region, angs, mags, used, drawImg);
        int numCells = (int)regStats[0];
	double orientation = regStats[1];
        LSR newLSR (numCells, region, orientation, img_gray);
	if(newLSR.isKeeper()){
	    allLSRs[LSRcount] = newLSR;
            LSRcount++;
	} else {
            for(int j=0;j<numCells;j++){
                int rejCell = region[j];
                used[rejCell] = 0;
                if(SHOULD_DRAW){
		    int *pix = addrToCoords(rejCell);
	    	    Point pt1 = Point(pix[0],pix[1]);
		    line(drawImg, pt1, pt1, Scalar(150, 150, 0), 1);
		}
            }
            used[nextPix] = 1; //don't want to try it again
        }
	if(SHOULD_DRAW){
	    imshow("Growing region...", drawImg);
	    waitKey(0);
        }
    }
    return LSRcount;
}

/**
*  Given an unfinished LSR, computes approximation of the best fit line.
*/
int *findLine(int *bb, int center, double orient){
    int *linePts = new int[2];
    int *lineMax = new int[2];
    int *lineMin = new int[2];
    int *centCoords = addrToCoords(center);
    bool vertical = 0;
    if(orient < (3./4)*PI && orient >= (1./4)*PI) vertical = 1;
    if(vertical){ //vertically-inclined
        int upIters = bb[3] - centCoords[1]; //dist between centroid and upper lim
        int downIters = centCoords[1] - bb[1]; //dist between centroid and lower lim
        double dx = 1/tan(orient); //if we treat dy as 1, this is dx
        lineMax[0] = max(centCoords[0] - dx*upIters,0.0);
        lineMax[1] = bb[3];
        lineMin[0] = centCoords[0] + dx*downIters;
        lineMin[1] = bb[1];
    } else { //horizontally-inclined
        int rightIters = bb[2] - centCoords[0]; //dist between centroid and right lim
        int leftIters = centCoords[0] - bb[0]; //dist between centroid and left lim
        double dy = tan(orient); //if we treat dx as 1, this is dy
        lineMax[0] = bb[2];
        lineMax[1] = max(centCoords[1] - dy*rightIters,0.0);  //added 'max' stuff to prevent endpoints < 0
        lineMin[0] = bb[0];
        lineMin[1] = centCoords[1] + dy*leftIters;
    }

    linePts[0] = coordsToAddr(lineMin[0],lineMin[1]);
    linePts[1] = coordsToAddr(lineMax[0],lineMax[1]);
    return linePts;
}

//////////////// Neighbors \\\\\\\\\\\\\\\\\\\\

/**
*  Returns a list of cells in the neighborhood of cellAddr,
*  Starting from the immediate neighbors and working outwards.
*  radius refers to the number of 'rings' around the center we'll 
*  look at (use r=1 for only immediate neighbors).
*/
int *findNeighborhood(int cellAddr, int radius){    
    int *coords = addrToCoords(cellAddr);
    int xCoord = coords[0];
    int yCoord = coords[1];
    int neighborhood_width = 2*radius + 1;
    int neighborhood_size = neighborhood_width*neighborhood_width - 1;
    int *neighborhood = new int[neighborhood_size];
    int index = 0;
    for(int r = 1; r<=radius; r++){
        //top and bottom of ring
        for(int j = yCoord-r; j <= yCoord+r; j += 2*r){
            for(int i = xCoord-r; i <= xCoord+r; i++){
                if(i<0 || i>=XDIM || j<0 || j>=YDIM)
                    neighborhood[index] = -1;
                else
                    neighborhood[index] = coordsToAddr(i, j);
                ++index;
            }
        }
        //sides of ring
        for(int i = xCoord-r; i <= xCoord+r; i += 2*r){
            for(int j = yCoord-r+1; j < yCoord+r; j++){
                if(i<0 || i>=XDIM || j<0 || j>=YDIM)
                    neighborhood[index] = -1;
                else
                    neighborhood[index] = coordsToAddr(i, j);
                ++index;
            }
        }
    }
    return neighborhood;
}











